package com.huawei.M01.iterator.M17;

import java.util.HashMap;
import java.util.Map;

public class SubArrays02 {
    public int countSubarrays(int[] nums, int k) {
        int n = nums.length;
        int kIndex = -1;
        for (int i = 0; i < n; i++) {
            if (nums[i] == k) {
                kIndex = i;
                break;
            }
        }
        // 统计个数
        int res = 0;
        Map<Integer, Integer> cache = new HashMap<>();
        int sum = 0; // 前缀和
        cache.put(0, 1);
        for (int i = 0; i < n; i++) {
            sum += gen(nums[i] - k);
            if (i < kIndex) {
                cache.put(sum, cache.getOrDefault(sum, 0) + 1);
            } else {
                Integer pre1 = cache.getOrDefault(sum, 0);
                Integer pre2 = cache.getOrDefault(sum - 1, 0);
                res += pre1 + pre2;
            }
        }
        return res;
    }

    private int gen(int num) {
        if (num == 0) {
            return 0;
        }
        return num > 0 ? 1 : -1;
    }
}
